import { Injectable, NestMiddleware } from '@nestjs/common';
import { Request, Response, NextFunction } from 'express';
import { v4 as uuidv4 } from 'uuid';
import { setRequestContext, RequestContext } from '@/common/async-context';

@Injectable()
export class TraceIdMiddleware implements NestMiddleware {
    use(req: Request, res: Response, next: NextFunction) {
        // 从请求头获取 traceId，如果没有则生成新的
        const xTraceId = 'x-trace-id';
        const traceId = (req.headers[xTraceId] as string) || uuidv4().replaceAll('-', '');
        const uri = req.url;

        // 创建请求上下文
        const requestContext: RequestContext = {
            traceId,
            uri,
            timestamp: new Date().toISOString(),
        };

        // 设置到请求对象中（保持向后兼容）
        (req as any).traceId = traceId;
        (req as any).uri = uri;

        // 设置到响应头中
        res.setHeader(xTraceId, traceId);

        // 使用 AsyncLocalStorage 设置上下文
        setRequestContext(requestContext);
        next();
    }
}
